{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Randomness and PRNGs in Flax"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this guide, you will learn how Flax uses [JAX's explicit pseudorandom number generator (PRNG) keys](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#jax-prng) to emulate randomness, and adds some additional features to make it easier for users to thread PRNG keys through different Flax `Module`s.\n",
    "\n",
    "If you are new to JAX PRNG keys or need a refresher, check out:\n",
    "- [JAX 101: PRNGs in JAX](https://jax.readthedocs.io/en/latest/jax-101/05-random-numbers.html)\n",
    "- [🔪 JAX - The Sharp Bits 🔪: Random Numbers](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#random-numbers)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install or upgrade Flax, and then import some necessary dependencies.\n",
    "\n",
    "**Note:** This guide uses the `--xla_force_host_platform_device_count=8` flag to emulate multiple devices in a CPU environment in a Google Colab/Jupyter Notebook. You don’t need this if you are already using a multi-device Google Cloud TPU environment, for example, on Google Cloud or in a Kaggle VM with a TPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": [
     "skip-execution"
    ]
   },
   "outputs": [],
   "source": [
    "!pip install -q flax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"XLA_FLAGS\"] = '--xla_force_host_platform_device_count=8'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import flax, flax.linen as nn\n",
    "import jax, jax.numpy as jnp\n",
    "from jax.sharding import Mesh, PartitionSpec, NamedSharding\n",
    "from jax.experimental import mesh_utils\n",
    "from jax.experimental.shard_map import shard_map\n",
    "\n",
    "import hashlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "outputId": "ec904f6b-0e87-4efe-87c4-fea0f8e8ec23"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[CpuDevice(id=0),\n",
       " CpuDevice(id=1),\n",
       " CpuDevice(id=2),\n",
       " CpuDevice(id=3),\n",
       " CpuDevice(id=4),\n",
       " CpuDevice(id=5),\n",
       " CpuDevice(id=6),\n",
       " CpuDevice(id=7)]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.devices()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the JAX config variable `jax_threefry_partitionable` to `True`. This will be the default value in the future and makes the PRNG more efficiently auto-parallelizable under `jax.jit`. Refer to [JAX discussion](https://github.com/jax-ml/jax/discussions/18480) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "jax.config.update('jax_threefry_partitionable', True)\n",
    "assert jax.config.jax_threefry_partitionable == True\n",
    "assert jax.config.jax_default_prng_impl == 'threefry2x32'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Receiving, manipulating and creating PRNG keys with `Module.make_rng`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The primary method Flax uses to receive, manipulate and create PRNG keys is via the `Module` method [`self.make_rng`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.make_rng). It is a method that accepts a string name that represents an \"RNG stream\". Each RNG stream has an initial starting seed PRNG key, which the user passes in as a dictionary argument (i.e. into an [`.init`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.init) or [`.apply`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.apply) function), and the starting seed is used by `self.make_rng` to generate more PRNG keys for that stream. If `self.make_rng` is called on a string name that does not have an initial starting seed PRNG key (i.e. the user did not pass a key with the corresponding name into `.init` or `.apply`), then `self.make_rng` will use the `'params'` key as the initial starting seed by default.\n",
    "\n",
    "Note that this method can only be called with bounded modules (see [The Flax Module lifecycle](https://flax.readthedocs.io/en/latest/developer_notes/module_lifecycle.html#top-level-modules))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "outputId": "2a16435b-e92a-480a-f9fb-e6effc42c4c2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Array((), dtype=key<fry>) overlaying:\n",
      "[1428664606 3351135085]\n",
      "Array((), dtype=key<fry>) overlaying:\n",
      "[3456700291 3873160899]\n",
      "Array((), dtype=key<fry>) overlaying:\n",
      "[2411773124 4124888837]\n"
     ]
    }
   ],
   "source": [
    "class RNGModule(nn.Module):\n",
    "  @nn.compact\n",
    "  def __call__(self):\n",
    "    print(self.make_rng('rng_stream'))\n",
    "    print(self.make_rng('rng_stream'))\n",
    "    print(self.make_rng('rng_stream'))\n",
    "\n",
    "rng_module = RNGModule()\n",
    "variables = rng_module.init({'rng_stream': jax.random.key(0)})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now if we use a different starting seed PRNG key, we will generate different values (as intended)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "outputId": "985b0f62-dfde-4f0f-fad4-a31927fc9f59"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Array((), dtype=key<fry>) overlaying:\n",
      "[3077990774 2166202870]\n",
      "Array((), dtype=key<fry>) overlaying:\n",
      "[3825832496 2886313970]\n",
      "Array((), dtype=key<fry>) overlaying:\n",
      "[ 791337683 1373966058]\n"
     ]
    }
   ],
   "source": [
    "variables = rng_module.init({'rng_stream': jax.random.key(1)})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calling `self.make_rng` for one stream will not affect the random values generated from another stream; i.e. the call order doesn't matter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "outputId": "7e8ce538-e380-4db9-db23-bc4a8da577da"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rng_stream1: Array((), dtype=key<fry>) overlaying:\n",
      "[1428664606 3351135085]\n",
      "rng_stream2: Array((), dtype=key<fry>) overlaying:\n",
      "[3077990774 2166202870]\n",
      "rng_stream1: Array((), dtype=key<fry>) overlaying:\n",
      "[3456700291 3873160899]\n",
      "rng_stream2: Array((), dtype=key<fry>) overlaying:\n",
      "[3825832496 2886313970]\n",
      "rng_stream1: Array((), dtype=key<fry>) overlaying:\n",
      "[2411773124 4124888837]\n",
      "rng_stream2: Array((), dtype=key<fry>) overlaying:\n",
      "[ 791337683 1373966058]\n"
     ]
    }
   ],
   "source": [
    "class RNGModuleTwoStreams(nn.Module):\n",
    "  @nn.compact\n",
    "  def __call__(self):\n",
    "    # same value as first code snippet above\n",
    "    print(f\"rng_stream1: {self.make_rng('rng_stream1')}\")\n",
    "    # same value as second code snippet above\n",
    "    print(f\"rng_stream2: {self.make_rng('rng_stream2')}\")\n",
    "    # same value as first code snippet above\n",
    "    print(f\"rng_stream1: {self.make_rng('rng_stream1')}\")\n",
    "    # same value as second code snippet above\n",
    "    print(f\"rng_stream2: {self.make_rng('rng_stream2')}\")\n",
    "    # same value as first code snippet above\n",
    "    print(f\"rng_stream1: {self.make_rng('rng_stream1')}\")\n",
    "    # same value as second code snippet above\n",
    "    print(f\"rng_stream2: {self.make_rng('rng_stream2')}\")\n",
    "\n",
    "rng_module_two_streams = RNGModuleTwoStreams()\n",
    "variables = rng_module_two_streams.init(\n",
    "  {'rng_stream1': jax.random.key(0), 'rng_stream2': jax.random.key(1)}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Providing the same seed PRNG key will result in the same values being generated (provided that the same operations are used for those keys)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "outputId": "b70be039-589a-48f7-dc54-65e78c449c65"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rng_stream1: Array((), dtype=key<fry>) overlaying:\n",
      "[1428664606 3351135085]\n",
      "rng_stream2: Array((), dtype=key<fry>) overlaying:\n",
      "[1428664606 3351135085]\n",
      "rng_stream1: Array((), dtype=key<fry>) overlaying:\n",
      "[3456700291 3873160899]\n",
      "rng_stream2: Array((), dtype=key<fry>) overlaying:\n",
      "[3456700291 3873160899]\n",
      "rng_stream1: Array((), dtype=key<fry>) overlaying:\n",
      "[2411773124 4124888837]\n",
      "rng_stream2: Array((), dtype=key<fry>) overlaying:\n",
      "[2411773124 4124888837]\n"
     ]
    }
   ],
   "source": [
    "variables = rng_module_two_streams.init(\n",
    "  {'rng_stream1': jax.random.key(0), 'rng_stream2': jax.random.key(0)}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### How `self.make_rng` works under the hood"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is what happens when `self.make_rng` (`flax.linen.Module.make_rng`) is called:\n",
    "* The following data is collected:\n",
    "  * The path of the `Module` as provided by `self.scope.path` (the top-level root module has an empty path `()`).\n",
    "  * The `self.make_rng` call count. That is, the number of times `self.make_rng` has been called for this specific stream (including this call).\n",
    "    * **Note:** Each sub-`Module` will have its own individual call count that's separate from other `Module`s. For example, a `Module` that has called `self.make_rng('params')` twice and contains a sub-`Module` that has called `self.make_rng('params')` once, will have a call count of 2 and 1 for each of the RNG stream `'params'`, respectively.\n",
    "* The data is bundled into a tuple and fed into a hash function and produces an integer.\n",
    "* The generated integer is folded into the RNG stream's starting seed PRNG key to generate a new, unique PRNG key."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below is a slightly simplified version of the hash function that Flax uses for `self.make_rng`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def produce_hash(data):\n",
    "  m = hashlib.sha1()\n",
    "  for x in data:\n",
    "    if isinstance(x, str):\n",
    "      m.update(x.encode('utf-8'))\n",
    "    elif isinstance(x, int):\n",
    "      m.update(x.to_bytes((x.bit_length() + 7) // 8, byteorder='big'))\n",
    "    else:\n",
    "      raise ValueError(f'Expected int or string, got: {x}')\n",
    "  d = m.digest()\n",
    "  hash_int = int.from_bytes(d[:4], byteorder='big')\n",
    "  return hash_int"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now you can manually reproduce the PRNG keys generated from `self.make_rng`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "outputId": "d26b7355-9e8b-4954-b2f4-cf7520d5c5a3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Array((), dtype=key<fry>) overlaying:\n",
      "[1428664606 3351135085]\n",
      "Array((), dtype=key<fry>) overlaying:\n",
      "[3456700291 3873160899]\n",
      "Array((), dtype=key<fry>) overlaying:\n",
      "[2411773124 4124888837]\n"
     ]
    }
   ],
   "source": [
    "stream_seed = jax.random.key(0)\n",
    "for call_count in range(1, 4):\n",
    "  hash_int = produce_hash(data=(call_count,))\n",
    "  print(jax.random.fold_in(stream_seed, jnp.uint32(hash_int)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "outputId": "dec627a6-4c5a-4e3e-ce11-ce4f72775261"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Array((), dtype=key<fry>) overlaying:\n",
      "[1428664606 3351135085]\n",
      "Array((), dtype=key<fry>) overlaying:\n",
      "[3456700291 3873160899]\n",
      "Array((), dtype=key<fry>) overlaying:\n",
      "[2411773124 4124888837]\n"
     ]
    }
   ],
   "source": [
    "variables = rng_module.init({'rng_stream': jax.random.key(0)})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sub-`Module`s and `self.make_rng`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This section explores how `self.make_rng` (`flax.linen.Module.make_rng`) behaves with sub-`Module`s.\n",
    "\n",
    "Consider the following example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "outputId": "5b7a9ae9-ca49-4ac0-d007-5caeee739ff0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RNGModule, count 1: Array((), dtype=key<fry>) overlaying:\n",
      "[1428664606 3351135085]\n",
      "RNGModule, count 2: Array((), dtype=key<fry>) overlaying:\n",
      "[3456700291 3873160899]\n",
      "RNGSubModule_0, count 1: Array((), dtype=key<fry>) overlaying:\n",
      "[3858825717 2323087578]\n",
      "RNGSubModule_0, count 2: Array((), dtype=key<fry>) overlaying:\n",
      "[ 601859108 3782857444]\n",
      "RNGSubSubModule_0, count 1: Array((), dtype=key<fry>) overlaying:\n",
      "[ 234240654 1028548813]\n",
      "RNGSubSubModule_0, count 2: Array((), dtype=key<fry>) overlaying:\n",
      "[3650462303 2124609379]\n"
     ]
    }
   ],
   "source": [
    "class RNGSubSubModule(nn.Module):\n",
    "  def __call__(self):\n",
    "    print(f\"{self.name}, count 1: {self.make_rng('rng_stream')}\")\n",
    "    print(f\"{self.name}, count 2: {self.make_rng('rng_stream')}\")\n",
    "\n",
    "class RNGSubModule(nn.Module):\n",
    "  @nn.compact\n",
    "  def __call__(self):\n",
    "    print(f\"{self.name}, count 1: {self.make_rng('rng_stream')}\")\n",
    "    print(f\"{self.name}, count 2: {self.make_rng('rng_stream')}\")\n",
    "    RNGSubSubModule()()\n",
    "\n",
    "class RNGModule(nn.Module):\n",
    "  @nn.compact\n",
    "  def __call__(self):\n",
    "    print(f\"RNGModule, count 1: {self.make_rng('rng_stream')}\")\n",
    "    print(f\"RNGModule, count 2: {self.make_rng('rng_stream')}\")\n",
    "    RNGSubModule()()\n",
    "\n",
    "rng_module = RNGModule()\n",
    "variables = rng_module.init({'rng_stream': jax.random.key(0)})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As previously discussed, the data that is fed into the Flax hash function consists of:\n",
    "\n",
    "  * The path of the `Module`, provided by `self.scope.path` (the top-level root module has an empty path `()`); and\n",
    "  * The call count for the specific RNG stream.\n",
    "\n",
    "In addition, note that each Flax `Module` and sub-`Module` have their own individual call counts, even for the same RNG stream. The convention for sub-`Module` names is: `f'{module_name}_{module_number}'`. For example, the first `Dense` sub-`Module` will be called `Dense_0`, the second one will be called `Dense_1`, and so on.\n",
    "\n",
    "Therefore, the following data will be fed into the hash function:\n",
    "\n",
    "  * For `RNGModule`: The data is just the call count, such as `(1,)` and `(2,)`, since the root `Module` has an empty path.\n",
    "  * For `RNGSubModule`: The data is `('RNGSubModule_0', 1)` and `('RNGSubModule_0', 2)`.\n",
    "  * For `RNGSubSubModule`: The data is `('RNGSubModule_0', 'RNGSubSubModule_0', 1)` and `('RNGSubModule_0', 'RNGSubSubModule_0', 2)`.\n",
    "\n",
    "With this data, you can manually reproduce the PRNG keys generated from the `Module` and sub-`Module`s using `self.make_rng`.\n",
    "\n",
    "For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "outputId": "c0de4d37-0f00-4e58-bdfd-e8a6454ed681"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RNGModule, count 1: Array((), dtype=key<fry>) overlaying:\n",
      "[1428664606 3351135085]\n",
      "RNGModule, count 2: Array((), dtype=key<fry>) overlaying:\n",
      "[3456700291 3873160899]\n",
      "RNGSubModule_0, count 1: Array((), dtype=key<fry>) overlaying:\n",
      "[3858825717 2323087578]\n",
      "RNGSubModule_0, count 2: Array((), dtype=key<fry>) overlaying:\n",
      "[ 601859108 3782857444]\n",
      "RNGSubSubModule_0, count 1: Array((), dtype=key<fry>) overlaying:\n",
      "[ 234240654 1028548813]\n",
      "RNGSubSubModule_0, count 2: Array((), dtype=key<fry>) overlaying:\n",
      "[3650462303 2124609379]\n"
     ]
    }
   ],
   "source": [
    "stream_seed = jax.random.key(0)\n",
    "for initial_data in ((), ('RNGSubModule_0',), ('RNGSubModule_0', 'RNGSubSubModule_0')):\n",
    "  if initial_data:\n",
    "    module_name = initial_data[-1]\n",
    "  else:\n",
    "    module_name = 'RNGModule'\n",
    "  for call_count in (1, 2):\n",
    "    hash_int = produce_hash(data=initial_data+(call_count,))\n",
    "    rng_key = jax.random.fold_in(stream_seed, jnp.uint32(hash_int))\n",
    "    print(f\"{module_name}, count {call_count}: {rng_key}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the same sub-`Module` class is used multiple times, you can increment the suffix of the sub-`Module` name accordingly. For example: `RNGSubModule_0`, `RNGSubModule_1`, and so on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "outputId": "0b77a038-7000-407b-c5b8-a28dea7951d1"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RNGModule, count 1: Array((), dtype=key<fry>) overlaying:\n",
      "[1428664606 3351135085]\n",
      "RNGModule, count 2: Array((), dtype=key<fry>) overlaying:\n",
      "[3456700291 3873160899]\n",
      "RNGSubModule_0, count 1: Array((), dtype=key<fry>) overlaying:\n",
      "[3858825717 2323087578]\n",
      "RNGSubModule_0, count 2: Array((), dtype=key<fry>) overlaying:\n",
      "[ 601859108 3782857444]\n",
      "RNGSubModule_1, count 1: Array((), dtype=key<fry>) overlaying:\n",
      "[ 426957352 2006350344]\n",
      "RNGSubModule_1, count 2: Array((), dtype=key<fry>) overlaying:\n",
      "[4006253729 4205356731]\n"
     ]
    }
   ],
   "source": [
    "class RNGSubModule(nn.Module):\n",
    "  @nn.compact\n",
    "  def __call__(self):\n",
    "    print(f\"{self.name}, count 1: {self.make_rng('rng_stream')}\")\n",
    "    print(f\"{self.name}, count 2: {self.make_rng('rng_stream')}\")\n",
    "\n",
    "class RNGModule(nn.Module):\n",
    "  @nn.compact\n",
    "  def __call__(self):\n",
    "    print(f\"RNGModule, count 1: {self.make_rng('rng_stream')}\")\n",
    "    print(f\"RNGModule, count 2: {self.make_rng('rng_stream')}\")\n",
    "    RNGSubModule()()\n",
    "    RNGSubModule()()\n",
    "\n",
    "rng_module = RNGModule()\n",
    "variables = rng_module.init({'rng_stream': jax.random.key(0)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "outputId": "d189d25e-425d-4fd7-fe18-2dfd63f28b87"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RNGModule, count 1: Array((), dtype=key<fry>) overlaying:\n",
      "[1428664606 3351135085]\n",
      "RNGModule, count 2: Array((), dtype=key<fry>) overlaying:\n",
      "[3456700291 3873160899]\n",
      "RNGSubModule_0, count 1: Array((), dtype=key<fry>) overlaying:\n",
      "[3858825717 2323087578]\n",
      "RNGSubModule_0, count 2: Array((), dtype=key<fry>) overlaying:\n",
      "[ 601859108 3782857444]\n",
      "RNGSubModule_1, count 1: Array((), dtype=key<fry>) overlaying:\n",
      "[ 426957352 2006350344]\n",
      "RNGSubModule_1, count 2: Array((), dtype=key<fry>) overlaying:\n",
      "[4006253729 4205356731]\n"
     ]
    }
   ],
   "source": [
    "stream_seed = jax.random.key(0)\n",
    "for initial_data in ((), ('RNGSubModule_0',), ('RNGSubModule_1',)):\n",
    "  if initial_data:\n",
    "    module_name = initial_data[-1]\n",
    "  else:\n",
    "    module_name = 'RNGModule'\n",
    "  for call_count in (1, 2):\n",
    "    hash_int = produce_hash(data=initial_data+(call_count,))\n",
    "    rng_key = jax.random.fold_in(stream_seed, jnp.uint32(hash_int))\n",
    "    print(f\"{module_name}, count {call_count}: {rng_key}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using `self.param` and `self.variable`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Flax users have the option of creating additional parameters and variables in their modules by using the [`self.param`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.param) and [`self.variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.variable) `Module` methods. An `init_fn` argument must be passed to these methods so that it can generate the initial value of the parameter/variable. `self.make_rng` is commonly used implicitly or explicitly in this `init_fn`, since many initializer functions are stochastic in nature and require a PRNG key. See the full list of Flax initializers [here](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/initializers.html).\n",
    "\n",
    "There are a couple of differences between the two methods that the user should take note of:\n",
    "* `self.param` always creates a parameter in the `'params'` [collection](https://flax.readthedocs.io/en/latest/glossary.html#term-Params-parameters), whereas `self.variable` creates a variable in any [collection](https://flax.readthedocs.io/en/latest/glossary.html#term-Variable-collections) the user specifies\n",
    "* `self.param` will automatically call `self.make_rng('params')` and pass in the generated PRNG key implicitly to the `init_fn` of the parameter you instantiated (it will be passed in as the first argument), whereas users will have to manually specify what RNG stream to call `self.make_rng` on in the `init_fn` of `self.variable` (it could be `'params'` or something different).\n",
    "\n",
    "Below is an example using both `self.param` and `self.variable`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "outputId": "a7816385-0e08-48e2-dc51-055d7bcd0bab"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-1.6185919   0.700908  ]\n",
      " [-1.3146383  -0.79342234]]\n",
      "[[ 0.0761425 -1.6157459]\n",
      " [-1.6857724  0.7126891]]\n",
      "[[ 0.60175574  0.2553228 ]\n",
      " [ 0.27367848 -2.1975214 ]]\n",
      "[[1.6249592  0.30813068]\n",
      " [1.6613585  1.0404155 ]]\n",
      "[[ 0.0030665   0.29551846]\n",
      " [ 0.16670242 -0.78252524]]\n",
      "[1.582462   0.15216611]\n"
     ]
    }
   ],
   "source": [
    "class Model(nn.Module):\n",
    "  @nn.compact\n",
    "  def __call__(self, x):\n",
    "    # kernel will use 'params' seed, initial data will include 'Dense_0', call count 1\n",
    "    x = nn.Dense(2, kernel_init=jax.random.normal, use_bias=False)(x)\n",
    "    # model_param will use 'params' seed, call count 1\n",
    "    model_param = self.param('model_param', jax.random.normal, x.shape)\n",
    "    # model_variable1 will use 'params' seed, call count 2\n",
    "    model_variable1 = self.variable(\n",
    "      'other_collection',\n",
    "      'model_variable1',\n",
    "      lambda: jax.random.normal(self.make_rng('params'), x.shape),\n",
    "    )\n",
    "    # model_variable2 will use 'other' seed, call count 1\n",
    "    model_variable2 = self.variable(\n",
    "      'other_collection',\n",
    "      'model_variable2',\n",
    "      lambda: jax.random.normal(self.make_rng('other'), x.shape),\n",
    "    )\n",
    "    # kernel will use 'params' seed, initial data will include 'Dense_1', call count 1\n",
    "    # bias will use 'params' seed, initial data will include 'Dense_1', call count 2\n",
    "    x = nn.Dense(2, kernel_init=jax.random.normal, bias_init=jax.random.normal)(\n",
    "      x\n",
    "    )\n",
    "    return x\n",
    "\n",
    "model = Model()\n",
    "variables = model.init(\n",
    "  {'params': jax.random.key(0), 'other': jax.random.key(1)}, jnp.ones((2, 2))\n",
    ")\n",
    "print(variables['params']['Dense_0']['kernel'])\n",
    "print(variables['params']['model_param'])\n",
    "print(variables['other_collection']['model_variable1'])\n",
    "print(variables['other_collection']['model_variable2'])\n",
    "print(variables['params']['Dense_1']['kernel'])\n",
    "print(variables['params']['Dense_1']['bias'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remember:\n",
    "* there is a separate count for each RNG stream; this is why the count for `self.make_rng('other')` starts at 1 even though there were earlier calls of `self.make_rng('params')`\n",
    "* each submodule has their own separate count for each rng stream; this is why each `Dense` layer has their own separate count for `self.make_rng('params')` and why `model_param` and `model_variable1` share the same count (since they are defined within the same top-level parent module)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "outputId": "ccec9d64-9a27-47f7-adaf-b36a5ea655db"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-1.6185919   0.700908  ]\n",
      " [-1.3146383  -0.79342234]]\n",
      "[[ 0.0761425 -1.6157459]\n",
      " [-1.6857724  0.7126891]]\n",
      "[[ 0.60175574  0.2553228 ]\n",
      " [ 0.27367848 -2.1975214 ]]\n",
      "[[1.6249592  0.30813068]\n",
      " [1.6613585  1.0404155 ]]\n",
      "[[ 0.0030665   0.29551846]\n",
      " [ 0.16670242 -0.78252524]]\n",
      "[[1.582462   0.15216611]]\n"
     ]
    }
   ],
   "source": [
    "params_seed = jax.random.key(0)\n",
    "other_seed = jax.random.key(1)\n",
    "for initial_data, count, seed, shape in (\n",
    "  (('Dense_0',), 1, params_seed, (2, 2)),\n",
    "  ((), 1, params_seed, (2, 2)),\n",
    "  ((), 2, params_seed, (2, 2)),\n",
    "  ((), 1, other_seed, (2, 2)),\n",
    "  (('Dense_1',), 1, params_seed, (2, 2)),\n",
    "  (('Dense_1',), 2, params_seed, (1, 2)),\n",
    "):\n",
    "  hash_int = produce_hash(data=(*initial_data, count))\n",
    "  rng_key = jax.random.fold_in(seed, jnp.uint32(hash_int))\n",
    "  print(jax.random.normal(rng_key, shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Managing RNG streams inside a training loop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below is an example of managing RNG streams from `self.make_rng`, `self.param`, `self.variable` and `nn.Dropout` in a training loop (note: `nn.Dropout` requires a seed PRNG key to be passed in the `'dropout'` RNG stream, since it implicitly calls `self.make_rng('dropout')`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SubModule(nn.Module):\n",
    "  @nn.compact\n",
    "  def __call__(self, x, train):\n",
    "    # variables created using `self.param` will use `self.make_rng('params')`\n",
    "    kernel = self.param('submodule_kernel', jax.random.normal, x.shape)\n",
    "    x = x + kernel\n",
    "    # `nn.Dropout` will use self.make_rng('dropout')\n",
    "    x = nn.Dropout(0.2)(x, deterministic=not train)\n",
    "    # `nn.Dense` will use self.make_rng('params')\n",
    "    x = nn.Dense(3)(x)\n",
    "    return x\n",
    "\n",
    "class Model(nn.Module):\n",
    "  @nn.compact\n",
    "  def __call__(self, x, train):\n",
    "    # make kernel use `self.make_rng('other')`\n",
    "    kernel = self.variable(\n",
    "      'other_collection',\n",
    "      'module_kernel',\n",
    "      lambda: jax.random.normal(self.make_rng('other'), x.shape),\n",
    "    )\n",
    "    x = (\n",
    "      x + kernel.value\n",
    "    )  # `.value` will extract the underlying value of the variable\n",
    "    x = SubModule()(x, train)\n",
    "    # `nn.Dropout` will use self.make_rng('dropout')\n",
    "    x = nn.Dropout(0.2)(x, deterministic=not train)\n",
    "    # `nn.Dense` will use self.make_rng('params')\n",
    "    x = nn.Dense(2)(x)\n",
    "    return x\n",
    "\n",
    "params_rng, other_rng, train_rng = jax.random.split(jax.random.key(0), 3)\n",
    "init_rngs = {'params': params_rng, 'other': other_rng}\n",
    "\n",
    "x = jnp.ones((1, 3))\n",
    "y = jnp.ones((1, 2))\n",
    "\n",
    "module = Model()\n",
    "variables = module.init(init_rngs, x, train=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "outputId": "e9da8228-acba-403d-bcb5-33a39d4d530d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.518454\n",
      "2.4859657\n",
      "2.4171872\n",
      "2.412684\n",
      "2.3435805\n",
      "2.2773488\n",
      "2.2592616\n",
      "2.2009292\n",
      "2.1839895\n",
      "2.1707344\n"
     ]
    }
   ],
   "source": [
    "def update(variables, rng):\n",
    "  # we don't need to provide a 'params' or 'other' rng, as only 'dropout' rng will be used during training\n",
    "  # split the rng to get a dropout_rng to be used for this training iteration,\n",
    "  # and to get another rng key to be used for the next training iteration\n",
    "  dropout_rng, next_rng = jax.random.split(rng)\n",
    "  def loss(params):\n",
    "    out = module.apply(\n",
    "      {'params': params, 'other_collection': variables['other_collection']},\n",
    "      x,\n",
    "      train=True,\n",
    "      rngs={'dropout': dropout_rng},\n",
    "    )\n",
    "    return jnp.mean((y - out) ** 2)\n",
    "  grads = jax.grad(loss)(variables['params'])\n",
    "  params = jax.tree_util.tree_map(lambda p, g: p - 1e-3 * g, variables['params'], grads)\n",
    "  return {\n",
    "    'params': params,\n",
    "    'other_collection': variables['other_collection'],\n",
    "  }, next_rng\n",
    "\n",
    "for _ in range(10):\n",
    "  variables, train_rng = update(variables, train_rng)\n",
    "  out = module.apply(variables, x, train=False)\n",
    "  print(jnp.mean((y - out)**2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🔪 Sharp edge 🔪 - unintentionally generating the same values\n",
    "\n",
    "There is an edge case where the same value can be unintentionally generated.\n",
    "See the [Flax issue](https://github.com/google/flax/issues/2157) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "outputId": "887142ff-c9ca-4aae-d9fa-cc9993d809c5"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(True, dtype=bool)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Leaf(nn.Module):\n",
    "  def __call__(self, x):\n",
    "    return x + jax.random.randint(self.make_rng(\"rng\"), (), 0, 100)\n",
    "\n",
    "class Node(nn.Module):\n",
    "  leaf_name: str\n",
    "  @nn.compact\n",
    "  def __call__(self, x):\n",
    "    return Leaf(name=self.leaf_name)(x)\n",
    "\n",
    "class Model(nn.Module):\n",
    "  @nn.compact\n",
    "  def __call__(self, x):\n",
    "    return (Node(name=\"ab\", leaf_name=\"cdef\")(x),\n",
    "            Node(name=\"abc\", leaf_name=\"def\")(x),\n",
    "    )\n",
    "\n",
    "out1, out2 = Model().apply({}, 0, rngs={\"rng\": jax.random.key(33)})\n",
    "out1 == out2 # same output, despite having different submodule names"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This occurs because the hash function [concatenates strings together](https://docs.python.org/3/library/hashlib.html#hashlib.hash.update), so the data `('AB', 'C')` is equivalent to data `('A', 'BC')` when fed into the hash function, therefore producing the same hash int."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "outputId": "001cbd49-129b-4474-c6a1-3255a4ee3dfe"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "947574064\n",
      "947574064\n",
      "947574064\n",
      "947574064\n"
     ]
    }
   ],
   "source": [
    "print(produce_hash(data=('A', 'B', 'C', 1)))\n",
    "print(produce_hash(data=('AB', 'C', 1)))\n",
    "print(produce_hash(data=('A', 'BC', 1)))\n",
    "print(produce_hash(data=('ABC', 1)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To avoid this edge case, users can flip the `flax_fix_rng_separator` [configuration flag](https://flax.readthedocs.io/en/latest/api_reference/flax.config.html#flax.configurations.Config.flax_fix_rng_separator) to `True`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "outputId": "35a2b204-bdfd-4f83-8e98-ba723963cb0c"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(False, dtype=bool)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "flax.config.update('flax_fix_rng_separator', True)\n",
    "out1, out2 = Model().apply({}, 0, rngs={\"rng\": jax.random.key(33)})\n",
    "out1 == out2 # different output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Managing RNG's on multiple devices"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This section will show examples on how to use `jit` and `shard_map` to use RNG's in multi-device settings."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using `jax.jit`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When using [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html), we can use RNG's as we did before, but we now include `in_shardings` and `out_shardings` arguments to specify how to shard input and output data. The RNG key itself gets replicated (not sharded); `jax.jit` makes each device use it as appropriate for its shard of the data.\n",
    "\n",
    "For more details on training on multiple devices in Flax using `jax.jit`, see our [Scale up Flax Modules on multiple devices guide](https://flax.readthedocs.io/en/latest/guides/parallel_training/flax_on_pjit.html#) and [lm1b example](https://github.com/google/flax/tree/main/examples/lm1b)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "outputId": "6c280522-4b43-4b82-f40a-b73986659b2c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[CpuDevice(id=0) CpuDevice(id=1) CpuDevice(id=2) CpuDevice(id=3)\n",
      " CpuDevice(id=4) CpuDevice(id=5) CpuDevice(id=6) CpuDevice(id=7)]\n",
      "Mesh(device_ids=array([0, 1, 2, 3, 4, 5, 6, 7]), axis_names=('data',))\n",
      "NamedSharding(mesh=Mesh('data': 8), spec=PartitionSpec('data',))\n"
     ]
    }
   ],
   "source": [
    "# Create a mesh and annotate the axis with a name.\n",
    "device_mesh = mesh_utils.create_device_mesh((8,))\n",
    "print(device_mesh)\n",
    "\n",
    "mesh = Mesh(devices=device_mesh, axis_names=('data',))\n",
    "print(mesh)\n",
    "\n",
    "data_sharding = NamedSharding(mesh, PartitionSpec('data',))\n",
    "print(data_sharding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "outputId": "d1bfbcad-e28a-4fae-8136-98bd1efb9332"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[-2.2187614 ],\n",
       "       [-2.8055234 ],\n",
       "       [-2.5464187 ],\n",
       "       [ 1.0270392 ],\n",
       "       [-3.5243359 ],\n",
       "       [-2.2795477 ],\n",
       "       [-0.6504516 ],\n",
       "       [ 0.17373264]], dtype=float32)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Model(nn.Module):\n",
    "  @nn.compact\n",
    "  def __call__(self, x, add_noise):\n",
    "    x = nn.Dense(1)(x)\n",
    "    # use jnp.where for control flow; for more details see: https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError\n",
    "    return jnp.where(\n",
    "      add_noise, x + jax.random.normal(self.make_rng('params'), x.shape), x\n",
    "    )\n",
    "\n",
    "module = Model()\n",
    "init_rng, apply_rng = jax.random.split(jax.random.key(0))\n",
    "x = jnp.ones((8, 1))\n",
    "variables = module.init(init_rng, x, False)\n",
    "\n",
    "# create custom forward function, since jit does not support kwargs when in_shardings is specified\n",
    "def forward(variables, x, add_noise, rng):\n",
    "  return module.apply(variables, x, add_noise, rngs={'params': rng})\n",
    "\n",
    "# shard the inputs x across devices\n",
    "# replicate the variables, add_noise boolean and rng key across devices\n",
    "# shard the output across devices\n",
    "jit_forward = jax.jit(\n",
    "  forward,\n",
    "  in_shardings=(None, data_sharding, None, None),\n",
    "  out_shardings=data_sharding,\n",
    ")\n",
    "out = jit_forward(variables, x, True, apply_rng)\n",
    "out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output is different given the same input, meaning the RNG key was used to add noise to the output.\n",
    "\n",
    "We can also confirm that the output is sharded across devices:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "outputId": "b672b85f-7a2d-44b5-afc1-bbf9426655ed"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Shard(device=CpuDevice(id=0), index=(slice(0, 1, None), slice(None, None, None)), replica_id=0, data=[[-2.2187614]]),\n",
       " Shard(device=CpuDevice(id=1), index=(slice(1, 2, None), slice(None, None, None)), replica_id=0, data=[[-2.8055234]]),\n",
       " Shard(device=CpuDevice(id=2), index=(slice(2, 3, None), slice(None, None, None)), replica_id=0, data=[[-2.5464187]]),\n",
       " Shard(device=CpuDevice(id=3), index=(slice(3, 4, None), slice(None, None, None)), replica_id=0, data=[[1.0270392]]),\n",
       " Shard(device=CpuDevice(id=4), index=(slice(4, 5, None), slice(None, None, None)), replica_id=0, data=[[-3.5243359]]),\n",
       " Shard(device=CpuDevice(id=5), index=(slice(5, 6, None), slice(None, None, None)), replica_id=0, data=[[-2.2795477]]),\n",
       " Shard(device=CpuDevice(id=6), index=(slice(6, 7, None), slice(None, None, None)), replica_id=0, data=[[-0.6504516]]),\n",
       " Shard(device=CpuDevice(id=7), index=(slice(7, 8, None), slice(None, None, None)), replica_id=0, data=[[0.17373264]])]"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out.addressable_shards"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another way to visualize the output sharding:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "outputId": "1c0a16ce-fa3f-4b95-d794-58464bbaa9ae"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  CPU 0  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">  CPU 1  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">  CPU 2  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  CPU 3  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">  CPU 4  </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">  CPU 5  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">  CPU 6  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">  CPU 7  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">         </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;214;97;107m  \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107mCPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;214;97;107m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;140;162;82m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82mCPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;140;162;82m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mCPU 3\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m  \u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148mCPU 4\u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148m  \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;107;110;207m  \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207mCPU 5\u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;107;110;207m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;165;81;148m  \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148mCPU 6\u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;165;81;148m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;140;109;49m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49mCPU 7\u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;140;109;49m         \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "jax.debug.visualize_array_sharding(out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we choose not to add noise, then the output is the same across all batches (as expected, since the input is the same for all batches):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "outputId": "fe9ec875-3e7f-4861-babc-f07064737276"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[-1.2839764],\n",
       "       [-1.2839764],\n",
       "       [-1.2839764],\n",
       "       [-1.2839764],\n",
       "       [-1.2839764],\n",
       "       [-1.2839764],\n",
       "       [-1.2839764],\n",
       "       [-1.2839764]], dtype=float32)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out = jit_forward(variables, x, False, apply_rng)\n",
    "out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can confirm the un-jitted function produces the same values, albeit unsharded (note there may be small numerical differences due to compiler optimizations from jitting):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "outputId": "0a9e5f2c-d4bf-4051-bf71-f32a9c32dc06"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[-2.2187614 ],\n",
       "       [-2.8055234 ],\n",
       "       [-2.5464187 ],\n",
       "       [ 1.0270392 ],\n",
       "       [-3.5243359 ],\n",
       "       [-2.2795477 ],\n",
       "       [-0.6504516 ],\n",
       "       [ 0.17373264]], dtype=float32)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out = forward(variables, x, True, apply_rng)\n",
    "out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "outputId": "772a5063-1bd5-46b4-f6f6-cae9b4b81a26"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[-1.2839764],\n",
       "       [-1.2839764],\n",
       "       [-1.2839764],\n",
       "       [-1.2839764],\n",
       "       [-1.2839764],\n",
       "       [-1.2839764],\n",
       "       [-1.2839764],\n",
       "       [-1.2839764]], dtype=float32)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out = forward(variables, x, False, apply_rng)\n",
    "out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using `shard_map`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When using [`jax.experimental.shard_map.shard_map`](https://jax.readthedocs.io/en/latest/jep/14273-shard-map.html), the important parts to remember are to:\n",
    "* split your PRNG key to produce a different key for each device\n",
    "* the PRNG keys will be sharded automatically to each device (provided you use the correct partition specification), but the [**rank of the original batched PRNG key array will not be reduced**](https://jax.readthedocs.io/en/latest/jep/14273-shard-map.html#rank-reducing-vs-rank-preserving-maps-over-array-axes); e.g.\n",
    "with a batch of 8 PRNG keys and 8 devices, each device will see a PRNG key batch of size 1 within the `shard_map`-ed function\n",
    "  * therefore to access the PRNG key itself, we need to index slice into it (see the example below)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "outputId": "aa00b9a3-24ba-4048-ed8c-afbb9070f039"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[-1.2605132 ],\n",
       "       [-1.2405176 ],\n",
       "       [-0.99350417],\n",
       "       [-1.0277128 ],\n",
       "       [-1.4154483 ],\n",
       "       [-0.3905797 ],\n",
       "       [-2.417677  ],\n",
       "       [ 0.9023453 ]], dtype=float32)"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def forward(variables, x, add_noise, rng_key_batch):\n",
    "  # rng_key_batch is a batch of size 1 containing 1 PRNG key\n",
    "  # index slice into the rng_key_batch to access the PRNG key\n",
    "  return module.apply(\n",
    "    variables, x, add_noise, rngs={'params': rng_key_batch[0]}\n",
    "  )\n",
    "\n",
    "# define partition specifications\n",
    "data_pspec = PartitionSpec('data')\n",
    "no_pspec = PartitionSpec()\n",
    "\n",
    "# shard the inputs x and rng keys across devices\n",
    "# replicate the variables and add_noise boolean across devices\n",
    "# shard the output across devices\n",
    "shmap_forward = shard_map(\n",
    "  forward,\n",
    "  mesh=mesh,\n",
    "  in_specs=(no_pspec, data_pspec, no_pspec, data_pspec),\n",
    "  out_specs=data_pspec,\n",
    ")\n",
    "# get 8 different rng's that will be used by the 8 devices when doing forward inference\n",
    "apply_rngs = jax.random.split(apply_rng, 8)\n",
    "out = shmap_forward(variables, x, True, apply_rngs)\n",
    "out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Confirm that the output is sharded across devices:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "outputId": "e304289b-ef1c-4e4a-d4c1-4c41613bfa62"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Shard(device=CpuDevice(id=0), index=(slice(0, 1, None), slice(None, None, None)), replica_id=0, data=[[-1.2605132]]),\n",
       " Shard(device=CpuDevice(id=1), index=(slice(1, 2, None), slice(None, None, None)), replica_id=0, data=[[-1.2405176]]),\n",
       " Shard(device=CpuDevice(id=2), index=(slice(2, 3, None), slice(None, None, None)), replica_id=0, data=[[-0.99350417]]),\n",
       " Shard(device=CpuDevice(id=3), index=(slice(3, 4, None), slice(None, None, None)), replica_id=0, data=[[-1.0277128]]),\n",
       " Shard(device=CpuDevice(id=4), index=(slice(4, 5, None), slice(None, None, None)), replica_id=0, data=[[-1.4154483]]),\n",
       " Shard(device=CpuDevice(id=5), index=(slice(5, 6, None), slice(None, None, None)), replica_id=0, data=[[-0.3905797]]),\n",
       " Shard(device=CpuDevice(id=6), index=(slice(6, 7, None), slice(None, None, None)), replica_id=0, data=[[-2.417677]]),\n",
       " Shard(device=CpuDevice(id=7), index=(slice(7, 8, None), slice(None, None, None)), replica_id=0, data=[[0.9023453]])]"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out.addressable_shards"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "outputId": "52fdb6d2-4c4f-44b3-feee-4bc5363c8f2f"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  CPU 0  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">  CPU 1  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #d6616b\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">  CPU 2  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8ca252\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  CPU 3  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">  CPU 4  </span>\n",
       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #e7cb94\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">  CPU 5  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #6b6ecf\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">  CPU 6  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #a55194\">         </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">  CPU 7  </span>\n",
       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #8c6d31\">         </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;214;97;107m  \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107mCPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;214;97;107m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;140;162;82m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82mCPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;140;162;82m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mCPU 3\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m  \u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148mCPU 4\u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148m  \u001b[0m\n",
       "\u001b[38;2;0;0;0;48;2;231;203;148m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;107;110;207m  \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207mCPU 5\u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;107;110;207m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;165;81;148m  \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148mCPU 6\u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;165;81;148m         \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;140;109;49m  \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49mCPU 7\u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m  \u001b[0m\n",
       "\u001b[38;2;255;255;255;48;2;140;109;49m         \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "jax.debug.visualize_array_sharding(out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lifted transforms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Flax lifted transforms](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/transformations.html) allow you to use [JAX transforms](https://github.com/jax-ml/jax#transformations) with `Module` arguments. This section will show you how to control how PRNG keys are split in Flax lifted transforms.\n",
    "\n",
    "Refer to [Lifted transformations](https://flax.readthedocs.io/en/latest/developer_notes/lift.html) for more detail."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `nn.vmap`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use [`nn.vmap`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.vmap.html) to create a batched `Dense` layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "outputId": "f0830f6b-659c-446f-c933-7b2a430f8004"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'params': {'bias': Array([0., 0.], dtype=float32),\n",
       "  'kernel': Array([[-1.2488099 , -0.6127134 ],\n",
       "         [-0.07084481,  0.60130936]], dtype=float32)}}"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = jnp.ones((3, 2))\n",
    "\n",
    "BatchDense = nn.vmap(\n",
    "    nn.Dense,\n",
    "    in_axes=0, out_axes=0,\n",
    "    variable_axes={'params': None},\n",
    "    split_rngs={'params': False})\n",
    "\n",
    "BatchDense(2).init(jax.random.key(0), x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By denoting `variable_axes={'params': 0}'`, we vectorize the `params` Arrays on the first axis. However the parameter values generated are all identical to each other:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "outputId": "eef5c0ca-f8d5-4f25-8ce6-9f2f60622daf"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'params': {'bias': Array([[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]], dtype=float32),\n",
       "  'kernel': Array([[[-1.2488099 , -0.6127134 ],\n",
       "          [-0.07084481,  0.60130936]],\n",
       "  \n",
       "         [[-1.2488099 , -0.6127134 ],\n",
       "          [-0.07084481,  0.60130936]],\n",
       "  \n",
       "         [[-1.2488099 , -0.6127134 ],\n",
       "          [-0.07084481,  0.60130936]]], dtype=float32)}}"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "BatchDense = nn.vmap(\n",
    "    nn.Dense,\n",
    "    in_axes=0, out_axes=0,\n",
    "    variable_axes={'params': 0},\n",
    "    split_rngs={'params': False})\n",
    "\n",
    "BatchDense(2).init(jax.random.key(0), x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we also make `split_rngs={'params': True}`, then the PRNG key we provide is split across the variable axis (in this case, the batch axis 0), and we can generate different parameters for each batch input:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "outputId": "275699c3-ba48-403e-877d-07b65981cff5"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'params': {'bias': Array([[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]], dtype=float32),\n",
       "  'kernel': Array([[[-0.2526208 , -0.15088455],\n",
       "          [-1.1987205 , -0.40843305]],\n",
       "  \n",
       "         [[-0.7064888 , -1.108805  ],\n",
       "          [-0.938775  ,  1.4812315 ]],\n",
       "  \n",
       "         [[-0.59468937, -0.2502723 ],\n",
       "          [-1.33515   ,  0.5067442 ]]], dtype=float32)}}"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "BatchDense = nn.vmap(\n",
    "    nn.Dense,\n",
    "    in_axes=0, out_axes=0,\n",
    "    variable_axes={'params': 0},\n",
    "    split_rngs={'params': True})\n",
    "\n",
    "BatchDense(2).init(jax.random.key(0), x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Adding a variable via `self.variable` is straightforward:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "outputId": "c11a80bc-d865-4e2e-e059-4d6bcea79e09"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'params': {'Dense_0': {'bias': Array([[0., 0.],\n",
       "          [0., 0.],\n",
       "          [0., 0.]], dtype=float32),\n",
       "   'kernel': Array([[[-0.9079084 ,  0.76390624],\n",
       "           [-0.01285526,  0.4320353 ]],\n",
       "   \n",
       "          [[ 0.12398645,  0.7884565 ],\n",
       "           [ 1.5344163 ,  1.3186085 ]],\n",
       "   \n",
       "          [[-0.44171348,  0.43430036],\n",
       "           [-0.40732604,  0.29774475]]], dtype=float32)}},\n",
       " 'other_collection': {'kernel': Array([[-0.8193048 ,  0.711106  ],\n",
       "         [-0.37802765, -0.66705877],\n",
       "         [-0.44808003,  0.93031347]], dtype=float32)}}"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Model(nn.Module):\n",
    "  @nn.compact\n",
    "  def __call__(self, x):\n",
    "    x = nn.Dense(2)(x)\n",
    "    kernel = self.variable(\n",
    "      'other_collection',\n",
    "      'kernel',\n",
    "      lambda: jax.random.normal(self.make_rng('other'), x.shape),\n",
    "    )\n",
    "    return x + kernel.value\n",
    "\n",
    "BatchModel = nn.vmap(\n",
    "  Model,\n",
    "  in_axes=0,\n",
    "  out_axes=0,\n",
    "  variable_axes={'params': 0, 'other_collection': 0},\n",
    "  split_rngs={'params': True, 'other': True},\n",
    ")\n",
    "\n",
    "BatchModel().init({'params': jax.random.key(0), 'other': jax.random.key(1)}, x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can control which RNG stream to split, for example, if we only wanted to split the `'params'` RNG stream, then the variables generated from `self.variable` will be the same for each batch input:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "outputId": "fb16619c-c975-497d-c867-6fd5143b4507"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'params': {'Dense_0': {'bias': Array([[0., 0.],\n",
       "          [0., 0.],\n",
       "          [0., 0.]], dtype=float32),\n",
       "   'kernel': Array([[[-0.9079084 ,  0.76390624],\n",
       "           [-0.01285526,  0.4320353 ]],\n",
       "   \n",
       "          [[ 0.12398645,  0.7884565 ],\n",
       "           [ 1.5344163 ,  1.3186085 ]],\n",
       "   \n",
       "          [[-0.44171348,  0.43430036],\n",
       "           [-0.40732604,  0.29774475]]], dtype=float32)}},\n",
       " 'other_collection': {'kernel': Array([[ 0.44956833, -1.1854612 ],\n",
       "         [ 0.44956833, -1.1854612 ],\n",
       "         [ 0.44956833, -1.1854612 ]], dtype=float32)}}"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "BatchModel = nn.vmap(\n",
    "    Model,\n",
    "    in_axes=0, out_axes=0,\n",
    "    variable_axes={'params': 0, 'other_collection': 0},\n",
    "    split_rngs={'params': True, 'other': False})\n",
    "\n",
    "BatchModel().init({'params': jax.random.key(0), 'other': jax.random.key(1)}, x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also control which parameters / variables should be generated for each batch input, for example, if we only wanted `'params'` to generate separate parameters for each batch input:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "outputId": "f3a17d59-6f75-4408-caba-5769d4589263"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'params': {'Dense_0': {'bias': Array([[0., 0.],\n",
       "          [0., 0.],\n",
       "          [0., 0.]], dtype=float32),\n",
       "   'kernel': Array([[[-0.9079084 ,  0.76390624],\n",
       "           [-0.01285526,  0.4320353 ]],\n",
       "   \n",
       "          [[ 0.12398645,  0.7884565 ],\n",
       "           [ 1.5344163 ,  1.3186085 ]],\n",
       "   \n",
       "          [[-0.44171348,  0.43430036],\n",
       "           [-0.40732604,  0.29774475]]], dtype=float32)}},\n",
       " 'other_collection': {'kernel': Array([ 0.44956833, -1.1854612 ], dtype=float32)}}"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "BatchModel = nn.vmap(\n",
    "    Model,\n",
    "    in_axes=0, out_axes=0,\n",
    "    variable_axes={'params': 0, 'other_collection': None},\n",
    "    split_rngs={'params': True, 'other': False})\n",
    "\n",
    "BatchModel().init({'params': jax.random.key(0), 'other': jax.random.key(1)}, x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `nn.scan`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use [`nn.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.scan.html) to create a scanned `Module` layer (this is useful for simplifying repetitively stacked submodules):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "outputId": "29d1863b-809f-42ce-894c-1b0810faa41e"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'params': {'Dense_0': {'bias': Array([[0., 0.],\n",
       "          [0., 0.],\n",
       "          [0., 0.]], dtype=float32),\n",
       "   'kernel': Array([[[-0.07838312, -0.7422982 ],\n",
       "           [ 0.87488323,  0.13773395]],\n",
       "   \n",
       "          [[ 0.97309333,  0.9087693 ],\n",
       "           [-0.12564984, -1.0920651 ]],\n",
       "   \n",
       "          [[-0.99055105,  1.1499453 ],\n",
       "           [-0.15721127, -0.62520015]]], dtype=float32)}}}"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = jnp.ones((3, 2))\n",
    "\n",
    "class ResidualMLPBlock(nn.Module):\n",
    "  @nn.compact\n",
    "  def __call__(self, x, _):\n",
    "    h = nn.Dense(features=2)(x)\n",
    "    h = nn.relu(h)\n",
    "    return x + h, None # return an empty carry\n",
    "\n",
    "ScanMLP = nn.scan(\n",
    "      ResidualMLPBlock, variable_axes={'params': 0},\n",
    "      variable_broadcast=False, split_rngs={'params': True},\n",
    "      length=3)\n",
    "\n",
    "ScanMLP().init(jax.random.key(0), x, None) # pass in an empty carry"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similar to before, we can control whether to split the RNG stream or not, for example, if we wanted all the stacked modules to be initialized to the same parameter values, we can pass in `split_rngs={'params': False}`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "outputId": "6a825bcd-9c3b-43c2-afd2-42500d89fb26"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'params': {'Dense_0': {'bias': Array([[0., 0.],\n",
       "          [0., 0.],\n",
       "          [0., 0.]], dtype=float32),\n",
       "   'kernel': Array([[[-0.66715515, -0.0484313 ],\n",
       "           [ 0.9867164 ,  0.75408363]],\n",
       "   \n",
       "          [[-0.66715515, -0.0484313 ],\n",
       "           [ 0.9867164 ,  0.75408363]],\n",
       "   \n",
       "          [[-0.66715515, -0.0484313 ],\n",
       "           [ 0.9867164 ,  0.75408363]]], dtype=float32)}}}"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ScanMLP = nn.scan(\n",
    "      ResidualMLPBlock, variable_axes={'params': 0},\n",
    "      variable_broadcast=False, split_rngs={'params': False},\n",
    "      length=3)\n",
    "\n",
    "ScanMLP().init(jax.random.key(0), x, None)"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,md:myst"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
